"""
性能测试工具模块
"""

import numpy as np
from datetime import datetime
from typing import List, Dict, Any
import logging
import sys
import os

# ========== 添加路径设置 ==========
# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, '..'))
if project_root not in sys.path:
    sys.path.insert(0, project_root)
# =================================

# 现在可以正常导入了
try:
    from utils.logger import setup_logger
except ImportError as e:
    print(f"导入错误: {e}")
    print("正在使用备用方案...")

    # 备用方案：直接在这里定义logger
    def setup_logger(name):
        import logging
        logger = logging.getLogger(name)
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(logging.INFO)
        return logger

logger = setup_logger(__name__)

class PerformanceTester:
    """性能测试工具类"""

    def __init__(self, searcher):
        self.searcher = searcher
        logger.info("PerformanceTester 初始化成功")

    def test_batch_sizes(self, batch_sizes: List[int] = None) -> List[Dict[str, Any]]:
        """测试不同batch_size对性能的影响"""
        if batch_sizes is None:
            batch_sizes = [1, 10, 100, 1000]

        # 检查是否有加载向量
        if self.searcher.name_vectors is None or len(self.searcher.ids) == 0:
            logger.warning("没有加载向量数据，使用模拟测试")
            return self._mock_batch_test(batch_sizes)

        results = []
        for batch_size in batch_sizes:
            # 生成测试数据
            vector_dim = self.searcher.name_vectors.shape[1]
            test_names = np.random.randn(batch_size, vector_dim)
            test_specs = np.random.randn(batch_size, vector_dim)
            test_units = [self.searcher.units[0] if self.searcher.units else "个"] * batch_size

            start_time = datetime.now()
            self.searcher.search_similar(test_names, test_specs, test_units, top_k=5)
            elapsed = (datetime.now() - start_time).total_seconds()

            results.append({
                'batch_size': batch_size,
                'time_seconds': elapsed,
                'qps': batch_size / elapsed if elapsed > 0 else 0
            })

            logger.info(f"Batch size {batch_size}: {elapsed:.3f}s, QPS: {batch_size/elapsed:.1f}")

        return results

    def test_data_sizes(self, data_sizes: List[int] = None) -> List[Dict[str, Any]]:
        """测试不同数据量下的性能"""
        if data_sizes is None:
            data_sizes = [1000, 10000, 50000, 100000]

        results = []
        for data_size in data_sizes:
            try:
                # 尝试导入NumpySearcher
                try:
                    from core.numpy_searcher import NumpySearcher
                except ImportError:
                    logger.warning("无法导入NumpySearcher，使用模拟测试")
                    return self._mock_data_test(data_sizes)

                # 生成测试数据
                vector_dim = 768  # 假设向量维度为768
                test_names = np.random.randn(10, vector_dim)  # 10个查询
                test_specs = np.random.randn(10, vector_dim)
                test_units = ["个"] * 10

                # 生成对应数据量的向量库
                fake_ids = list(range(data_size))
                fake_name_vecs = np.random.randn(data_size, vector_dim)
                fake_spec_vecs = np.random.randn(data_size, vector_dim)
                fake_units = ["个"] * data_size

                # 临时加载数据
                temp_searcher = NumpySearcher()
                temp_searcher.load_vectors(fake_ids, fake_name_vecs, fake_spec_vecs, fake_units)

                start_time = datetime.now()
                temp_searcher.search_similar(test_names, test_specs, test_units, top_k=5)
                elapsed = (datetime.now() - start_time).total_seconds()

                results.append({
                    'data_size': data_size,
                    'time_seconds': elapsed,
                    'avg_query_time': elapsed / 10  # 10个查询的平均时间
                })

                logger.info(f"Data size {data_size}: {elapsed:.3f}s, Avg: {elapsed/10:.4f}s/query")

            except Exception as e:
                logger.error(f"测试数据量 {data_size} 时出错: {e}")
                results.append({
                    'data_size': data_size,
                    'time_seconds': -1,
                    'avg_query_time': -1,
                    'error': str(e)
                })

        return results

    def _mock_batch_test(self, batch_sizes: List[int]) -> List[Dict[str, Any]]:
        """模拟批量测试（用于测试环境）"""
        logger.info("使用模拟批量测试")
        results = []
        for batch_size in batch_sizes:
            # 模拟性能数据
            time_seconds = max(0.001, batch_size * 0.0001)  # 模拟处理时间
            results.append({
                'batch_size': batch_size,
                'time_seconds': time_seconds,
                'qps': batch_size / time_seconds,
                'is_mock': True  # 标记为模拟数据
            })
        return results

    def _mock_data_test(self, data_sizes: List[int]) -> List[Dict[str, Any]]:
        """模拟数据量测试（用于测试环境）"""
        logger.info("使用模拟数据量测试")
        results = []
        for data_size in data_sizes:
            # 模拟性能数据
            time_seconds = max(0.1, data_size * 0.00001)  # 模拟处理时间
            results.append({
                'data_size': data_size,
                'time_seconds': time_seconds,
                'avg_query_time': time_seconds / 10,
                'is_mock': True  # 标记为模拟数据
            })
        return results

    def generate_report(self) -> Dict[str, Any]:
        """生成性能测试报告"""
        try:
            batch_results = self.test_batch_sizes()
            data_results = self.test_data_sizes()

            return {
                'batch_performance': batch_results,
                'data_size_performance': data_results,
                'test_time': datetime.now().isoformat(),
                'status': 'completed'
            }
        except Exception as e:
            logger.error(f"生成报告时出错: {e}")
            return {
                'status': 'error',
                'error': str(e),
                'test_time': datetime.now().isoformat()
            }

def create_test_data(num_samples: int, vector_dim: int = 768) -> tuple:
    """
    创建测试数据

    Args:
        num_samples: 样本数量
        vector_dim: 向量维度

    Returns:
        tuple: (ids, name_vectors, spec_vectors, units)
    """
    ids = list(range(num_samples))
    name_vectors = np.random.randn(num_samples, vector_dim)
    spec_vectors = np.random.randn(num_samples, vector_dim)
    units = ["个"] * num_samples

    return ids, name_vectors, spec_vectors, units

# 测试代码
if __name__ == "__main__":
    print("PerformanceTester 模块测试")

    # 创建测试搜索器
    class MockSearcher:
        def __init__(self):
            self.ids = []
            self.name_vectors = None
            self.spec_vectors = None
            self.units = []

        def load_vectors(self, ids, name_vectors, spec_vectors, units):
            self.ids = ids
            self.name_vectors = name_vectors
            self.spec_vectors = spec_vectors
            self.units = units

        def search_similar(self, query_names, query_specs, query_units, top_k=5):
            # 模拟搜索延迟
            import time
            time.sleep(0.001 * len(query_names))  # 模拟处理时间
            return [[] for _ in range(len(query_names))]

    # 创建测试数据
    mock_searcher = MockSearcher()
    test_ids, test_name_vecs, test_spec_vecs, test_units = create_test_data(1000)
    mock_searcher.load_vectors(test_ids, test_name_vecs, test_spec_vecs, test_units)

    # 测试性能测试器
    tester = PerformanceTester(mock_searcher)

    # 测试批量性能
    print("测试批量性能...")
    batch_results = tester.test_batch_sizes([1, 5, 10])
    for result in batch_results:
        print(f"Batch {result['batch_size']}: {result['time_seconds']:.3f}s")

    # 测试数据量性能
    print("\n测试数据量性能...")
    data_results = tester.test_data_sizes([100, 1000])
    for result in data_results:
        print(f"Data {result['data_size']}: {result['avg_query_time']:.4f}s/query")

    print("\n性能测试完成！")